# coding: UTF-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pdb
from sklearn.metrics import accuracy_score,precision_score,recall_score


class Config(object):

    """配置参数"""
    def __init__(self, n_vocab, num_classes):
        self.dropout = 0.5                                              # 随机失活
        self.require_improvement = 1000                                 # 若超过1000batch效果还没提升，则提前结束训练
        self.num_classes = num_classes                                  # 类别数
        self.n_vocab = n_vocab
        self.embed = 300
        self.filter_sizes = (2, 3, 4)                                   # 卷积核尺寸
        self.num_filters = 256                                          # 卷积核数量(channels数)


'''Convolutional Neural Networks for Sentence Classification'''


class Model(nn.Module):
    def __init__(self, config):
        super(Model, self).__init__()
        self.embedding = nn.Embedding(config.n_vocab, config.embed, padding_idx=config.n_vocab - 1)
        self.convs = nn.ModuleList(
            [nn.Conv2d(1, config.num_filters, (k, config.embed)) for k in config.filter_sizes])
        self.dropout = nn.Dropout(config.dropout)
        self.fc = nn.Linear(config.num_filters * len(config.filter_sizes), config.num_classes)
        self.criterion = nn.CrossEntropyLoss()

    def conv_and_pool(self, x, conv):
        x = F.relu(conv(x)).squeeze(3)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)
        return x

    def forward(self, x, label=None):
        embed = self.embedding(x)
        embed = embed.unsqueeze(1)
        conv = torch.cat([self.conv_and_pool(embed, conv) for conv in self.convs], 1)
        conv = self.dropout(conv)
        logit = self.fc(conv)
        pred = torch.argmax(logit,dim=1)
        if label is None:
            return pred
        else:
            loss = self.criterion(logit,label)
            acc = accuracy_score(label.tolist(),pred.tolist())
            precision = precision_score(label.tolist(),pred.tolist(),average="micro")
            recall = recall_score(label.tolist(),pred.tolist(),average="micro")
            return loss, acc, precision, recall

        